#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#define UP_DIV(x, y) (((x) + (y) - (1)) / (y))
#define PAD 1

__constant sampler_t
smp_none = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_NONE | CLK_FILTER_NEAREST;

constant FLT Bt[36] = {
        1.0000000000f, 0.0000000000f, -2.5000004768f, -0.0000001192f, 1.0000001192f, 0.0000000000f,
        0.0000000000f, 0.9428091049f, 1.3333333731f, -0.4714044929f, -0.6666667461f, 0.0000000000f,
        0.0000000000f, -0.9428089857f, 1.3333334923f, 0.4714045525f, -0.6666667461f, 0.0000000000f,
        0.0000000000f, -0.1178511307f, -0.0833333358f, 0.2357022613f, 0.1666666865f, 0.0000000000f,
        0.0000000000f, 0.1178511307f, -0.0833333507f, -0.2357022911f, 0.1666666865f, 0.0000000000f,
        0.0000000000f, 0.9999998808f, -0.0000000596f, -2.5000000000f, 0.0000000000f, 1.0000000000f,
};

__kernel void Winograd4x4To36(__read_only image2d_t input,
                              __write_only image2d_t output,
                              int4 input_shape,     // N IH IW CI_SLICES
                              int4 output_shape)    // N 36 TILE_XY CI_SLICES
{
    int tile_xy = get_global_id(0);
    int row = get_global_id(1);
    int slice = get_global_id(2);

    int TILE_XY = output_shape.z;
    int SLICES = input_shape.w;
    if (tile_xy >= TILE_XY || row >= 6 || slice >= SLICES)
    {
        return;
    }

    int IH = input_shape.y, IW = input_shape.z;
    int TILE_X = UP_DIV(IW, 4);
    int tile_x = tile_xy % TILE_X;
    int tile_y = tile_xy / TILE_X;

    constant FLT *Bt_row = Bt + row * 6;
    FLT4 BtD_row[6] = {0};
    for (int y = 0; y < 6; y++)
    {
        int ih = tile_y * 4 - PAD + y;

        // NC4HW4
        if(ih < 0 || ih >= IH) {continue;}
        int y_idx = slice * IH + ih;
        // NHWC4
//        int y_idx = ih;

        for (int x = 0; x < 6; x++)
        {
            int iw = tile_x * 4 - PAD + x;

            // NC4HW4
            int x_idx = iw;
            // NHWC4
//            if(iw < 0 || iw >= IW) {continue;}
//            int x_idx = iw * SLICES + slice;

            BtD_row[x] += Bt_row[y] * READ_IMAGE(input, smp_none, (int2)(x_idx, y_idx));
        }
    }

    for (int y = 0; y < 6; y++)
    {
        FLT4 acc = (FLT4)(0.0f, 0.0f, 0.0f, 0.0f);
        for (int x = 0; x < 6; x++)
        {
            acc += BtD_row[x] * Bt[y * 6 + x];
        }
        WRITE_IMAGE(output, (int2)(tile_xy, slice * 36 + (row * 6 + y)), acc); // NC4HW4: height=C*36 width=TILE_XY
    }
}